#!/bin/bash

# cd /home/chenzhongzhi/text-img-model
# DS_BUILD_FUSED_ADAM=1 pip install -U deepspeed==0.6.5
# export PYTHONPATH=$PYTHONPATH:`pwd`/src
export NCCL_SOCKET_IFNAME=eth0
export NCCL_IB_DISABLE=0
export NCCL_IB_CUDA_SUPPORT=1
export NCCL_IB_GID_INDEX=0
export NCCL_IB_HCA=mlx5_2,mlx5_5
export OMP_NUM_THREADS=4
echo "master_ip: $MASTER_ADDR"
# cd src
lr=2e-6
wd=2e-1
wu=2000
lg=laion2b_fromT
ep=10
echo $lg

# rm -rf logs/m9plus_v1_ft_${lg}_${lr}_${wd}_${wu}_$ep
# laion2b-multi
train_data="../../clip_contrast/src/datas/laion400m-dat-release/{000000..00300}.tar"

# laion400m-en
# train_data="laion400m/laion400m-full-release/img_data/laion400m-dat-release/{00000..41455}.tar"
# multilingual laion + laion6plus (M18)
# train_data="cl_dataset/m9_origin/m9_6plus_wudao_${lg}.json"
# text_ckpt="/share/project/text-img-model/xlmr_m9plus_v1.pt"
# text_ckpt="/home/chenzhongzhi/text-img-model/text-img-model/xlmr_9lg_18m_clip.pt"

# echo "${text_ckpt}"
TOKENIZERS_PARALLELISM=false  WANDB_MODE=online HF_DATASETS_OFFLINE=1 python3 -m torch.distributed.launch --nproc_per_node=8 \
        --nnodes=$WORLD_SIZE --node_rank=$RANK \
    --master_addr=$MASTER_ADDR --master_port=28375 --use_env \
    main_deepspeed.py \
        --dataset-type="webdataset" \
        --train-num-samples 10000000 \
        --train-data="${train_data}" \
        --dataset-resampled \
        --save-frequency 1 \
        --zeroshot-frequency 1 \
        --report-to="wandb" \
        --wandb-notes="" \
        # --imagenet-val="/share/projset/baaishare/baai-mrnd/datasets/imagenet2012/val/" \
        --warmup $wu \
        --batch-size=128 \
        --epochs=$ep \
        --lr=$lr \
        --wd=$wd \
        --norm_gradient_clip=5.0 \
        --workers=4 \
        --model ViT-L-14 \
        --name="m9plus_v1_ft_${lg}_${lr}_${wd}_${wu}_${ep}" \
        --seed 3407 \
        --text-distil='xlmr-large' \
        --gather-with-grad \
        --language-zh \
        --lock-image \
        # --text-ckpt "${text_ckpt}" \
        --env_type="pytorchDDP" 